`timescale 1ps/1ps
module aes_top(
    input   clk
,   input   rst_n

,   input   keyexp      
,   output  keyexp_done

,   input   encode
,   output  encode_done

,   input  [127:0] data
,   output [127:0] out

);

wire  [3:0]    key_addr;
wire  [127:0]  key_ex;
wire  [31:0]   key_sbox_addr;
wire  [31:0]   key_sbox_out;

wire  [2:0]    encode_sbox_cmd;
wire  [127:0]  encode_sbox_addr;
wire  [31:0]   encode_sbox_out_32;

wire [2:0]    sbox_cmd;
wire [127:0]  sbox_in;
wire [31:0]   sbox_out;

/* 密钥拓展状态 */ 
reg keyex_state;
always @(posedge clk or negedge rst_n) begin
    if(!rst_n)
        keyex_state <= 1'b0;
    else if(keyexp)
        keyex_state <= 1'b1;
    else if(keyexp_done)
        keyex_state <= 1'b0;
end

keyexp  u_keyexp (
    .clk                     ( clk         ),
    .rst_n                   ( rst_n       ),
    .en                      ( keyexp      ),//
    .addr                    ( key_addr    ),//
    .key                     ( data        ),//
    
    .done                    ( keyexp_done   ),//
    .exp_key                 ( key_ex        ),//
    .sbox_addr               ( key_sbox_addr ),//
    .sbox_out                ( key_sbox_out  ) //
);

encode u_encode (
    .clk                     ( clk           ),
    .rst_n                   ( rst_n         ),

    .en                      ( encode        ),//
    .txt                     ( data          ),//
    .key_in                  ( key_ex        ),//
    
    .done                    ( encode_done   ),//
    .out                     ( out           ),//
    .key_addr                ( key_addr      ),//
    .sbox_cmd                ( encode_sbox_cmd      ),//
    .sbox_out_32             ( encode_sbox_out_32   ),//
    .sbox_addr               ( encode_sbox_addr     ) // 
);


/* 复用sbox */
assign sbox_cmd = keyex_state ? 3'b111 : encode_sbox_cmd;
assign sbox_in  = keyex_state ? {96'd0,key_sbox_addr} : encode_sbox_addr;
assign encode_sbox_out_32 =  sbox_out;
assign key_sbox_out       =  sbox_out;

sbox4  share_sbox4 (
    .cmd                     ( sbox_cmd ),
    .in                      ( sbox_in  ),
    .out                     ( sbox_out )
);
endmodule